package com.by.log.filter;

import com.by.log.RpcLogContext;
import com.by.log.TraceContext;

import java.util.Map;
import java.util.UUID;

/**
 * RPC追踪过滤器
 * 用于在RPC调用中传递和管理追踪上下文信息
 */
public class RpcTraceFilter {
    private static final String TRACE_ID_KEY = "X-Trace-Id";
    private static final String SPAN_ID_KEY = "X-Span-Id";
    private static final String PARENT_SPAN_ID_KEY = "X-Parent-Span-Id";

    /**
     * 在RPC调用前设置追踪上下文
     * @param headers RPC调用的头部信息
     * @return RPC日志上下文
     */
    public static RpcLogContext beforeInvoke(Map<String, String> headers) {
        RpcLogContext context = new RpcLogContext();
        
        String traceId = headers.get(TRACE_ID_KEY);
        if (traceId == null || traceId.isEmpty()) {
            traceId = TraceContext.generateTraceId();
        }
        context.setTraceId(traceId);
        
        String parentSpanId = headers.get(PARENT_SPAN_ID_KEY);
        context.setParentSpanId(parentSpanId);
        
        String spanId = headers.get(SPAN_ID_KEY);
        if (spanId != null && !spanId.isEmpty()) {
            context.setSpanId(spanId);
        } else {
            context.setSpanId(TraceContext.generateSpanId());
        }
        
        // 设置当前线程的追踪上下文
        TraceContext traceContext = TraceContext.getCurrentContext();
        traceContext.setTraceId(traceId);
        traceContext.setSpanId(spanId);
        traceContext.setParentSpanId(parentSpanId);
        
        return context;
    }

    /**
     * 在RPC调用后清理追踪上下文
     */
    public static void afterInvoke() {
        TraceContext.clear();
    }
    
    /**
     * 获取当前追踪上下文，用于传递给下游服务
     * @return 包含追踪信息的头部映射
     */
    public static Map<String, String> getCurrentContextHeaders() {
        TraceContext traceContext = TraceContext.getCurrentContext();
        Map<String, String> headers = new java.util.HashMap<>();
        
        String traceId = traceContext.getTraceId();
        if (traceId != null) {
            headers.put(TRACE_ID_KEY, traceId);
        }
        
        String spanId = traceContext.getSpanId();
        if (spanId != null) {
            headers.put(SPAN_ID_KEY, spanId);
        }
        
        String parentSpanId = traceContext.getParentSpanId();
        if (parentSpanId != null) {
            headers.put(PARENT_SPAN_ID_KEY, parentSpanId);
        }
        
        return headers;
    }
    
    /**
     * 创建子跨度
     * @return 新的RPC日志上下文
     */
    public static RpcLogContext createChildSpan(String serviceName, String methodName) {
        TraceContext parentContext = TraceContext.getCurrentContext();
        
        RpcLogContext childContext = new RpcLogContext();
        childContext.setTraceId(parentContext.getTraceId());
        childContext.setParentSpanId(parentContext.getSpanId());
        childContext.setSpanId(TraceContext.generateSpanId());
        childContext.setServiceName(serviceName);
        childContext.setMethodName(methodName);
        
        // 更新当前线程的追踪上下文
        TraceContext currentContext = TraceContext.getCurrentContext();
        currentContext.setSpanId(childContext.getSpanId());
        currentContext.setParentSpanId(childContext.getParentSpanId());
        
        return childContext;
    }
}